import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    """
    一个简单的卷积神经网络模型，用于GTSRB数据集的交通标志分类。
    """
    def __init__(self, num_classes=43):
        super(SimpleCNN, self).__init__()
        # 卷积层 1: 输入通道=3 (RGB), 输出通道=32, 卷积核=5x5
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5)
        # 池化层 1: 2x2最大池化
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 卷积层 2: 输入通道=32, 输出通道=64, 卷积核=3x3
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        # 池化层 2: 2x2最大池化
        # self.pool (reused)
        
        # 卷积层 3: 输入通道=64, 输出通道=64, 卷积核=3x3
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3)
        # 池化层 3: 2x2最大池化
        # self.pool (reused)

        # 全连接层
        # The input size for the linear layer depends on the input image size.
        # Assuming input image is 32x32.
        # After conv1 (5x5): (32-5+1) = 28. After pool: 14x14.
        # After conv2 (3x3): (14-3+1) = 12. After pool: 6x6.
        # After conv3 (3x3): (6-3+1) = 4. After pool: 2x2.
        # Flattened size: 64 * 2 * 2 = 256
        self.fc1 = nn.Linear(in_features=64 * 2 * 2, out_features=128)
        self.fc2 = nn.Linear(in_features=128, out_features=num_classes)

        # Dropout layer
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        # 卷积和激活 -> 池化
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        
        # 展平特征图
        x = x.view(-1, 64 * 2 * 2)
        
        # 全连接层 -> Dropout -> 输出层
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

if __name__ == '__main__':
    # 测试模型是否能正常工作
    # 创建一个假的输入张量 (batch_size=4, channels=3, height=32, width=32)
    dummy_input = torch.randn(4, 3, 32, 32)
    # 实例化模型
    model = SimpleCNN(num_classes=43)
    print("模型结构:")
    print(model)
    # 前向传播
    output = model(dummy_input)
    print(f"\n输入尺寸: {dummy_input.shape}")
    print(f"输出尺寸: {output.shape}")
    assert output.shape == (4, 43)
    print("模型测试通过！") 